from yoloSeg import YoloSegPose
import torch
import numpy as np
from logger import setup_logger

logger = setup_logger()

class ComputePosition:
    def __init__(self):
        self.yoloSegPose = YoloSegPose()
        

    def get_pose_dict(self, result, poly, threshold=0.8):
        """
        return {
            'left_top': {'coord': (x, y)},
            'left_bottom': {'coord': (x, y)},
            'right_top': {'coord': (x, y)},
            'left_0': {'coord': (x, y)},
            'top_0': {'coord': (x, y)},
            ... # 其他点
        }
        poly: 多边形轮廓点集，形状为(N,2)的数组
        """
        # 确保多边形是二维数组 (n, 2)
        if poly.ndim == 3:
            poly = poly.reshape(-1, 2)
        
        # 1. 获取检测到的角点
        corner_names = ["left_top", "left_bottom", "right_top", "right_bottom"]
        corner_indices = [0, 1, 2, 3]
        
        xy = result.xy
        conf = result.conf
        if conf is None:
            print(f"未检测到关键点")
            return None
        
        if isinstance(xy, torch.Tensor):
            xy = xy.cpu().numpy()
        if isinstance(conf, torch.Tensor):
            conf = conf.cpu().numpy()

        pose_dict = {}
        for name, idx in zip(corner_names, corner_indices):
            x, y = xy[0][idx]
            c = conf[0][idx]
            if c >= threshold:
                pose_dict[name] = {"coord": (float(x), float(y))}
        
        # 2. 从多边形中移除已检测到的角点（最近点匹配）
        remaining_poly = poly.copy()  # 避免修改原始数据
        
        # 确保remaining_poly是二维数组
        if remaining_poly.ndim == 3:
            remaining_poly = remaining_poly.reshape(-1, 2)
        
        for name in list(pose_dict.keys()):
            if name not in corner_names:
                continue
            cx, cy = pose_dict[name]["coord"]
            # 找到最近的多边形点
            min_dist = float('inf')
            min_idx = -1
            for i, pt in enumerate(remaining_poly):
                px, py = pt
                dist = (px - cx) ** 2 + (py - cy) ** 2
                if dist < min_dist:
                    min_dist = dist
                    min_idx = i
            if min_idx >= 0:
                pose_dict[name] = {"coord": tuple(remaining_poly[min_idx])}
                remaining_poly = np.delete(remaining_poly, min_idx, axis=0)

        
        # 3. 根据检测情况推断边缘点
        detected_corners = set(pose_dict.keys())
        if detected_corners == {"left_top","left_bottom","right_top","right_bottom"} and len(remaining_poly) == 0:
            logger.info("检测到情况0: 完整漏出四个点")
        # 情况1: 只露出左上角 -> 三角形
        elif detected_corners == {"left_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况1: 只露出左上角，推断为三角形")
            # 按x坐标排序：左侧点(left_0)，右侧点(top_0)
            remaining_poly = sorted(remaining_poly, key=lambda p: p[0])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["top_0"] = {"coord": tuple(remaining_poly[1])}
            
        # 情况2: 露出左上+左下 -> 四边形
        elif detected_corners == {"left_top", "left_bottom"} and len(remaining_poly) == 2:
            logger.info("检测到情况2: 露出左上+左下，推断为四边形")
            # 按y坐标排序：上方点(top_0)，下方点(bottom_0)
            remaining_poly = sorted(remaining_poly, key=lambda p: p[1])
            pose_dict["top_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["bottom_0"] = {"coord": tuple(remaining_poly[1])}
        
        # 情况3: 露出左上+右上 -> 四边形
        elif detected_corners == {"left_top", "right_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况3: 露出左上+右上，推断为四边形")
            # 按x坐标排序：左侧点(left_0)，右侧点(right_0)
            remaining_poly = sorted(remaining_poly, key=lambda p: p[0])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[1])}
        
        # 情况4: 露出左上+左下+右上 -> 五边形
        elif detected_corners == {"left_top", "left_bottom", "right_top"} and len(remaining_poly) == 2:
            logger.info("检测到情况4: 露出左上+左下+右上，推断为五边形")
            # 按y坐标排序：上方点(right_0)，下方点(bottom_0)
            remaining_poly = sorted(remaining_poly, key=lambda p: p[1])
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["bottom_0"] = {"coord": tuple(remaining_poly[1])}
        
        # 情况5: 露出左下+右下 -> 四边形
        elif detected_corners == {"left_bottom", "right_bottom"} and len(remaining_poly) == 2:
            logger.info("检测到情况5: 露出左下+右下，推断为四边形")
            remaining_poly = sorted(remaining_poly, key=lambda p: p[0])
            pose_dict["left_0"] = {"coord": tuple(remaining_poly[0])}
            pose_dict["right_0"] = {"coord": tuple(remaining_poly[1])}

        logger.info(f"pose_dict: {pose_dict}")
        return pose_dict
    def get_polygon_pose(self, image, seg_results, pose_results):
        fitted_pose = {}
        num_pose_points=0
        if image is None:
            return num_pose_points,fitted_pose
        # 关键点检测
        if pose_results is None:
            # print("未检测到关键点")
            logger.info("关键点为None")
            return num_pose_points,fitted_pose
        if  len(pose_results.xy) == 0:
            # logger.info("pose.xy == 0:",pose_results.xy)
            return num_pose_points,fitted_pose
        # 分割检测
        if seg_results is None or len(seg_results) == 0 :
            print("未检测到分割结果")
            return num_pose_points,fitted_pose
        
        # 获取原始多边形
        polygon = seg_results
        # print("原始多边形",polygon)

        # 计算露出的角点数量（置信度大于0.5的角点）
        xy = pose_results.xy
        conf = pose_results.conf
        
        if conf is None:
            num_pose_points = 0
        else:
            if isinstance(xy, torch.Tensor):
                xy = xy.cpu().numpy()
            if isinstance(conf, torch.Tensor):
                conf = conf.cpu().numpy()
            
            # 只考虑前4个关键点（角点）
            corner_confs = conf[0][:4]
            num_pose_points = np.sum(corner_confs >= 0.8)
        
        # 多边形拟合
        fitted_polygon = self.yoloSegPose.get_polygon2(polygon, keypoints_num=num_pose_points)
        if fitted_polygon is None:
            print("多边形拟合失败")
            return num_pose_points,fitted_pose
        
        # 使用新的get_pose_dict函数获取完整的姿态点（包括推断的边缘点）
        pose_dict = self.get_pose_dict(pose_results, fitted_polygon)
        
        if pose_dict is None:
            return num_pose_points,fitted_pose
        
        # 直接使用get_pose_dict返回的点位
        fitted_pose = {name: info["coord"] for name, info in pose_dict.items()}
        return num_pose_points,fitted_pose

    def get_edges_3d_points(self,color_img,depth_img,seg_results,pose_results,visual_pose=False,
                            visual_edges=False,visual_edges_3d=False,
                            visual_fitted_edges=False,visual_fitted_vector=False):
        origin_img = color_img
        depth_img = depth_img
        # 1. 设置彩色图和深度图的大小
        self.yoloSegPose.calibration_tool.set_resolution(origin_img,depth_img)
        # 2. 设置深度图，因为后续都是使用同一个深度图，提前设置好,返回与Color对齐的深度图
        depth_img=self.yoloSegPose.calibration_tool.set_depth_img(depth_img)
        # 3. 获取角点 {"left_top":np.array,"left_bottom":np.array}
        pose_num,pose = self.get_polygon_pose(origin_img,seg_results,pose_results)
        if pose_num == 0 or pose is None or len(pose) == 0:
            logger.warning("角点数量为零,不应该在这里出现才对")
            return None,None
        else:
            logger.info(f"获取图片多边形关键点{pose_num}个,角点{len(pose)}个")

        if visual_pose:
            self.yoloSegPose.visualize_tool.visualize_pose(pose,origin_img)
        if len(pose) > 5:
            logger.warning("角点数量大于5,不可能出现才对")
            return None, None
        
        edges = self.yoloSegPose.processPoint.get_dict_edges(pose, 60)
        if edges is not None and len(edges) > 0:
            logger.info(f"获取图片多边形边{len(edges)}条")
        else:
            logger.warning("未获取到多边形边,不应该在这里出现才对")
            return None,None
        
        edges_depth = self.yoloSegPose.processPoint.get_dict_depth(edges)
        logger.info(f"获取图片多边形边深度{len(edges_depth)}个")
        if visual_edges:
            try:
                self.yoloSegPose.visualize_tool.visualize_edges(edges,edges_depth,origin_img,depth_img)
            except:
                print("visual_edges error:")
                print("edges:",edges)
                print("edges_depth:",edges_depth)
        edges_3d = self.yoloSegPose.processPoint.get_dict_3dpoints(edges,edges_depth)
        logger.info(f"计算图片多边形边3D点{len(edges_3d)}个")
        if visual_edges_3d:
            try:
                self.yoloSegPose.visualize_tool.visualize_edges_3d(edges_3d)
            except:
                print("visual_edges_3d error:")
                print(edges_3d)
        
        fitted_edges=self.yoloSegPose.processLine.get_dict_fitted_points(edges_3d)
        logger.info(f"计算图片多边形边拟合点{len(fitted_edges)}个")

        if visual_fitted_edges:
            try:
                self.yoloSegPose.visualize_tool.visualize_fitted_edges(edges_3d,fitted_edges)
            except:
                print("visual_fitted_edges error:")
                print("edges_3d:",edges_3d)
                print("fitted_edges:",fitted_edges)
        fitted_vector=self.yoloSegPose.processLine.get_dict_fitted_vector(fitted_edges)
        logger.info(f"计算图片多边形边拟合向量成功")
        if visual_fitted_vector:
            try:
                self.yoloSegPose.visualize_tool.visualize_fitted_vector(self.yoloSegPose.processVector.standard_vector,fitted_vector)
            except:
                print("visual_fitted_vector error:")
                print("standard_vector:",self.yoloSegPose.processVector.standard_vector)
                print("fitted_vector:",fitted_vector)
        return pose,fitted_vector

    def test_one_image(self, color_img, depth_img, seg_results, pose_results):
        if color_img is None:
            logger.warning("color_img is None")
            return None,None
        if depth_img is None:
            logger.warning("depth_img is None")
            return None,None
        if seg_results is None:
            logger.warning("seg_results is None")
            return None,None
        if pose_results is None:
            logger.warning("pose_results is None")
            return None,None
        pose,fitted_vector=self.get_edges_3d_points(color_img, depth_img, seg_results, pose_results)
        if pose is None:
            logger.warning("pose is None")
            return None,None
        if fitted_vector is None:
            logger.warning("fitted_vector is None")
            return None,None
        vector=self.yoloSegPose.processVector.get_xyzr_by_vector(fitted_vector)
        if vector is None:
            logger.warning("vector is None")
            return None,None
        return pose,vector

    def process_image_pair(self, timestamp, color_img, depth_img, seg_results, pose_results):
        pose,vector=self.test_one_image(color_img, depth_img, seg_results, pose_results)
        return pose,vector

    def set_standard(self, image_path,depth_path,seg_results,pose_results):
        pose,fitted_vector=self.get_edges_3d_points(image_path,depth_path,seg_results,pose_results)
        # vector=self.yoloSegPose.processVector.get_xyzr_by_vector(fitted_vector)
        if pose is not None and fitted_vector is not None:
            vector=self.yoloSegPose.processVector.set_standard(fitted_vector)
            vector=self.yoloSegPose.processVector.get_xyzr_by_vector(fitted_vector)
            return pose,vector
        else:
            return None,None

if __name__ == "__main__":
    compute_position = ComputePosition()
    
   
    